# -*- coding: utf-8 -*-
import random
import ConfigSpace as CS
import ConfigSpace.hyperparameters as CSH
import itertools

class AveragePool:
    cs_padding = CS.ConfigurationSpace()
    cs_attributes = CS.ConfigurationSpace()
    #attribute_list = ['auto_pad', 'ceil_mode', 'count_include_pad', 'kernel_shape', 'pads', 'strides']
    attribute_list = ['auto_pad', 'ceil_mode', 'count_include_pad', 'pads', 'strides']
    auto_pad_list = ['NOTSET', 'SAME_UPPER', 'SAME_LOWER', 'VALID']
    all_attributes = []

    def __init__(self):
        print('AveragePool Init--------------------------')
        self.init_all_attributes()
        self.init_cs()

    def init_all_attributes(self):
        for i in range(1, len(AveragePool.attribute_list)+1) :
            #此处取所有组合的可能值，存在列表里
            for attribute in itertools.combinations(AveragePool.attribute_list, i):
                l = list(attribute)
                AveragePool.all_attributes.append(l)

        #print('all_attributes: ', all_attributes)
        print('all attributes:')
        for a in AveragePool.all_attributes:
            print(a)

        tmp_attributes = []

        #将所有含有auto_pad的元素展开为'NOTSET', 'SAME_UPPER', 'SAME_LOWER', 'VALID'四种类型
        #如旧元素为['count_include_pad', 'auto_pad']，则展开后的元素为：
        #['count_include_pad', 'NOTSET']
        #['count_include_pad', 'SAME_UPPER']
        #['count_include_pad', 'SAME_LOWER']
        #['count_include_pad', 'VALID']

        for attribute in AveragePool.all_attributes :
            tmp_a = []

            if 'auto_pad' in attribute : 
                tmp = []
                for a in attribute :
                    if a != 'auto_pad' :
                        tmp.append(a)

                tmp1 = tmp.copy()
                tmp2 = tmp.copy() 
                tmp3 = tmp.copy() 
                tmp4 = tmp.copy() 

                tmp1.append('NOTSET') 
                tmp2.append('SAME_UPPER')
                tmp3.append('SAME_LOWER')       
                tmp4.append('VALID')

                tmp_a.append(tmp1)
                tmp_a.append(tmp2)
                tmp_a.append(tmp3)
                tmp_a.append(tmp4)

                tmp_attributes.append(tmp_a) 

        print('got new attributes: ')   
        for a in tmp_attributes:
                print(a) 

        for a in AveragePool.all_attributes[:]:
            if 'auto_pad' in a : 
                AveragePool.all_attributes.remove(a)

        #print('got last attributes: ') 
        for a in AveragePool.all_attributes:
            print(a)

        for a in tmp_attributes:
            AveragePool.all_attributes.extend(a)   

        for a in AveragePool.all_attributes:
            a.append('kernel_shape')

        l = ['kernel_shape']
        AveragePool.all_attributes.append(l)

        print('again, got last attributes: ') 
        for a in AveragePool.all_attributes:
            print(a)    
         
    def init_cs(self):
        padding = CSH.CategoricalHyperparameter('auto_pad', ['NOTSET', 'SAME_UPPER', 'SAME_LOWER', 'VALID'])
        AveragePool.cs_padding.add_hyperparameter(padding)

        length = len(AveragePool.all_attributes)

        attributes = CSH.UniformIntegerHyperparameter(name='combination', lower=0, upper=length-1)
        AveragePool.cs_attributes.add_hyperparameter(attributes)

    def get_random_padding(self):
        s = cs_padding.sample_configuration()
        padding = s.get_dictionary()['auto_pad']
        print('get auto_pad: ', padding)

    def get_random_combination(self):
        s = AveragePool.cs_attributes.sample_configuration()
        index = s.get_dictionary()['combination']
        print('get combination: ', index)

        return AveragePool.all_attributes[index]            

    def input_shape_correct(self, shapes):
        print('input_shape_correct, AveragePool')

        if len(shapes) != 1:
            return False

        if len(shapes[0]) < 3:
            for i in range(3 - len(shapes[0])):
                shapes[0].append(random.randint(2, 5))

        if len(shapes[0]) > 5:
            shapes[0] = shapes[0][:5]     

        return True    

    def get_max_attributes_combination(self):
        return len(AveragePool.all_attributes)

    def get_attributes_by_index_or_random(self, shapes, index):
        print('get_attributes_by_index_or_random, index: ', index)    
            
        attributes = {}
        shape_out = []

        stride_ = [1, 1, 1]
        pads_ = [0, 0, 0, 0, 0, 0]

        has_strides = False

        auto_pad_mode = 'NOTSET'

        if index != -1:
            attr_combination = AveragePool.all_attributes[index]
        else:
            attr_combination = self.get_random_combination()

        if len(attr_combination) == 0:
            return attributes, shape_out    

        kernel_shape = []

        if 'kernel_shape' in attr_combination:
            for i in range(len(shapes[0]) - 2):
                kernel_shape.append(random.randint(1, shapes[0][-1-i]))

            print('----got kernel_shape:', kernel_shape) 
            kernel_shape = kernel_shape[::-1]
            print('++++got kernel_shape:', kernel_shape) 

            attributes['kernel_shape'] = kernel_shape  

        if 'ceil_mode' in attr_combination:
            ceil_mode = random.randint(0, 1)
            attributes['ceil_mode'] = ceil_mode

        has_auto_pad = False        

        for attr in attr_combination:
            if attr in AveragePool.auto_pad_list:
                index = AveragePool.auto_pad_list.index(attr)
                print('set auto_pad: ', AveragePool.auto_pad_list[index])
                attributes['auto_pad'] = AveragePool.auto_pad_list[index]
                auto_pad_mode = AveragePool.auto_pad_list[index]
                has_auto_pad = True
                print('set auto_pad_mode: ', auto_pad_mode)

        if has_auto_pad == False:
            if 'pads' in attr_combination : 
                if len(shapes[0]) == 4:
                    pads = [random.randint(0, int((kernel_shape[0])/2)), random.randint(0, int((kernel_shape[1])/2)),
                                random.randint(0, int((kernel_shape[0])/2)), random.randint(0, int((kernel_shape[1])/2))]
                    attributes['pads'] = pads
                    pads_[0] = pads[0]
                    pads_[1] = pads[1]
                    pads_[2] = pads[2]
                    pads_[3] = pads[3]
                elif len(shapes[0]) == 5:
                    pads = [random.randint(0, int((kernel_shape[0])/2)), random.randint(0, int((kernel_shape[1])/2)),
                                random.randint(0, int((kernel_shape[2])/2)), random.randint(0, int((kernel_shape[0])/2)),
                                random.randint(0, int((kernel_shape[1])/2)), random.randint(0, int((kernel_shape[2])/2))]

                    attributes['pads'] = pads
                    pads_[0] = pads[0]
                    pads_[1] = pads[1]
                    pads_[2] = pads[2]
                    pads_[3] = pads[3]
                    pads_[4] = pads[4]
                    pads_[5] = pads[5]     
                else:
                    pads = [random.randint(0, int((kernel_shape[0])/2)), random.randint(0, int((kernel_shape[0])/2))]
                    attributes['pads'] = pads
                    pads_[0] = pads[0]
                    pads_[1] = pads[1]

        if 'strides' in attr_combination:
            if len(shapes[0]) == 4:
                strides = [random.randint(1, shapes[0][-2]), random.randint(1, shapes[0][-1])]
                attributes['strides'] = strides
                stride_[0] = strides[0]
                stride_[1] = strides[1]
            elif len(shapes[0]) == 5:
                strides = [random.randint(1, shapes[0][-3]), random.randint(1, shapes[0][-2]), random.randint(1, shapes[0][-1])]
                attributes['strides'] = strides
                stride_[0] = strides[0]
                stride_[1] = strides[1]
                stride_[1] = strides[2]     
            else:
                strides = [random.randint(1, shapes[0][-1])]
                attributes['strides'] = strides
                stride_[0] = strides[0] 

        return attributes, shape_out 


